import requests
import json
import base64


class Sonnet:
    base_url = "https://oa.api2d.net/claude/v1/messages"
    api_key = 'fk230351-xxx'


def encode_image_to_base64(image_path):
    with open(image_path, "rb") as image_file:
        return base64.b64encode(image_file.read()).decode('utf-8')


def chat_sonnet(prompt, model="claude-3-opus-20240229"):
    payload = json.dumps({
        "model": model,
        "messages": [
            {
                "role": "user",
                "content": prompt
            }
        ],
        "stream": False,
        "max_tokens": 100
    })
    headers = {
        'x-api-key': Sonnet.api_key,
        'Authorization': 'Bearer ' + Sonnet.api_key,
        'Content-Type': 'application/json'
    }
    response = requests.request("POST", Sonnet.base_url, headers=headers, data=payload)
    return response.json()


def img_chat_sonnet(prompt, image_path=None, base64_image='', model="claude-3-opus-20240229"):
    messages = [
        {
            "role": "user",
            "content": []
        }
    ]

    # 添加文本提示
    messages[0]["content"].append({
        "type": "text",
        "text": prompt
    })

    # 如果提供了图片，添加图片内容
    if image_path:
        base64_image = encode_image_to_base64(image_path)
        messages[0]["content"].append({
            "type": "image",
            "source": {
                "type": "base64",
                "media_type": "image/jpeg",
                "data": base64_image
            }
        })
    elif base64_image:
        messages[0]["content"].append({
            "type": "image",
            "source": {
                "type": "base64",
                "media_type": "image/jpeg",
                "data": base64_image
            }
        })

    payload = json.dumps({
        "model": model,
        "messages": messages,
        "stream": False,
        "max_tokens": 100
    })

    headers = {
        'x-api-key': Sonnet.api_key,
        'Authorization': 'Bearer ' + Sonnet.api_key,
        'Content-Type': 'application/json'
    }

    response = requests.request("POST", Sonnet.base_url, headers=headers, data=payload)
    return response.json()


if __name__ == '__main__':
    print(chat_sonnet('这是什么'))
    """
    {'id': 'msg_016dtVYE7mp8ou352WPNA7g4', 'type': 'message', 'role': 'assistant', 'model': 'claude-3-opus-20240229', 'content': [{'type': 'text', 'text': '好的,给你讲一个简短有趣的笑话:\n\n一天,小明在路上捡到了一个钱包。他打开一看,里面除了一些现金,还有一张纸条,上面写着:"我是隐形人,这个钱包是我不小心掉的,请还给我,我就在你旁边'}], 'stop_reason': 'max_tokens', 'stop_sequence': None, 'usage': {'input_tokens': 14, 'cache_creation_input_tokens': 0, 'cache_read_input_tokens': 0, 'output_tokens': 100, 'pre_token_count': 100, 'pre_total': 42, 'adjust_total': 3, 'final_total': 39}}
 
    """
    print(img_chat_sonnet('这是什么', r"C:\Users\HP\Pictures\yeye.jpeg"))
    """
    {'id': 'msg_01GqmT2uMt2oSjQUt3zAPpVs', 'type': 'message', 'role': 'assistant', 'model': 'claude-3-opus-20240229', 'content': [{'type': 'text', 'text': '这是一个3D渲染的老爷爷形象,由一个穿着蓝色针织开衫的卡通老人角色组成。他有一个白色的大胡子和眼镜,整体造型很可爱有童趣。从眼镜的镜片可以看出这是一个电脑生成的3D'}], 'stop_reason': 'max_tokens', 'stop_sequence': None, 'usage': {'input_tokens': 456, 'cache_creation_input_tokens': 0, 'cache_read_input_tokens': 0, 'output_tokens': 100, 'pre_token_count': 100, 'pre_total': 44, 'adjust_total': -28, 'final_total': 72}}
 
    """
